import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import argparse
from PIL import Image
from model import CNN
from prepare_data import *
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from math import sqrt
import warnings
warnings.filterwarnings('ignore')

def integers_to_onehot(integers, max_value=10):
    # 创建一个全为 0 的矩阵，行数为整数的个数，列数为 max_value + 1
    onehot_matrix = torch.zeros((len(integers), max_value + 1))
    # 为每个整数设置 one-hot 编码
    for i, integer in enumerate(integers):
        onehot_matrix[i, integer] = 1
    return onehot_matrix

def train(model,train_data):
    for epoch in range(opt.epoch):
        epoch_loss, iter = 0, 0
        for img, label in train_data:
            iter += 1
            img = img.to(device)
            # label = torch.tensor(label.unsqueeze(0), dtype=torch.float32).cuda()#work1
            label = torch.tensor(label.unsqueeze(0), dtype=torch.float32).cuda().long()#work2
            label = integers_to_onehot(label)
            print("label is", label)
            score = model(img)
            loss = loss_func(score, label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        epoch_loss /= iter
        print("loss is", epoch_loss)
    return model

def evaluate(score, label):
    score, label = torch.tensor(score), torch.tensor(label)
    mse = torch.mean(torch.abs(score - label) ** 2)
    mae = torch.mean(torch.abs(score - label))
    print(" mae:", mae, "mse:", mse)

def test(model,test_data):
    for epoch in range(opt.epoch):
        true_list, rec_list = [], []
        for img, label in test_data:
            img = img.to(device)
            label = label.to(device)
            score = model(img)
            # print("score",score)
            true_list.append(label.tolist())
            rec_list.append(score.tolist())
        evaluate(rec_list, true_list)

if __name__ == '__run2_':
    torch.cuda.set_device(0)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    SEED = 0
    torch.manual_seed(SEED)
    torch.cuda.manual_seed(SEED)

    parser = argparse.ArgumentParser()

    parser.add_argument('--batchSize', type=int, default=1, help='input batch size')  # 原本是10000 现改5000
    parser.add_argument('--hidden_size', type=int, default=1024, help='hidden state size')  # 原本是200 现改500
    parser.add_argument('--epoch', type=int, default=3, help='number of epochs to train for')
    parser.add_argument('--lr', type=float, default=0.001, help='learning rate')
    parser.add_argument('--l2', type=float, default=0.0001, help='l2 penalty')

    opt = parser.parse_args()

    train_set, test_set = get_data()
    train_data = DataLoader(dataset=train_set, batch_size=opt.batchSize, collate_fn=collate, shuffle=True,
                            pin_memory=True, num_workers=12)
    test_data = DataLoader(dataset=test_set, batch_size=opt.batchSize, collate_fn=collate, pin_memory=True,
                           num_workers=8)

    model = CNN().cuda()
    optimizer = optim.Adam(model.parameters(), lr=opt.lr, weight_decay=opt.l2)
    loss_func = nn.CrossEntropyLoss()
    # loss_func = nn.BCEWithLogitsLoss()
    model=train(model,train_data)
    test(model,train_data)